import backtrader as bt
import datetime
import utils
from strategy.utils.BaoStockPandasData import BaoStockPandasData
from utils.DataSource import get_stock_data2

def get_stock_data(stock_code, start_date, end_date):
    import baostock as bs
    import pandas as pd
    
    bs.login()
    rs = bs.query_history_k_data_plus(stock_code,
                                    fields="date,open,high,low,close,volume",
                                    start_date=start_date,
                                    end_date=end_date,
                                    adjustflag="3")
    data_list = []
    while (rs.error_code == '0') & rs.next():
        data_list.append(rs.get_row_data())
    bs.logout()
    df = pd.DataFrame(data_list, columns=rs.fields)
    df.open = df.open.astype(float)
    df.high = df.high.astype(float)
    df.low = df.low.astype(float)
    df.close = df.close.astype(float)
    df.volume = df.volume.astype(float)
    df.date= pd.to_datetime(df.date)
    return df

# 回测设置示例

class EMAStrategy(bt.Strategy):
    params = (
        ('short_window', 10),  # 短期EMA窗口大小
        ('long_window', 30),   # 长期EMA窗口大小
    )

    def __init__(self):
        # 初始化EMA指标
        self.short_ema = bt.indicators.ExponentialMovingAverage(
            self.data.close, period=self.params.short_window)
        self.long_ema = bt.indicators.ExponentialMovingAverage(
            self.data.close, period=self.params.long_window)

    def next(self):
        # 检查是否已经持仓
        if not self.position:
            # 如果短期EMA上穿长期EMA，买入
            if self.short_ema[0] > self.long_ema[0]:
                self.buy(size=100)
        else:
            # 如果短期EMA下穿长期EMA，卖出
            if self.short_ema[0] < self.long_ema[0]:
                self.sell(size=100)
# 回测设置示例
if __name__ == '__main__':
    cerebro = bt.Cerebro()

    # 假设已经通过某种方式（如Yahoo Finance）获取了数据
    # 假设df是从baostock获取并处理过的DataFrame
    df = get_stock_data2("sh.600000", "2020-01-01", "2023-12-31")
    data = BaoStockPandasData(dataname=df)
    cerebro.adddata(data)

    cerebro.addstrategy(EMAStrategy)

    cerebro.broker.setcash(100000.0)
    cerebro.broker.set_coc(True)  # 现金交易模式

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    cerebro.run()
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
    
#在Backtrader中实现均线收敛发散（Moving Average Convergence Divergence, MACD）策略，实际上是指利用MACD指标来进行交易决策，而非直接基于简单移动平均线（SMA）或指数移动平均线（EMA）的直接收敛与发散现象。然而，基于您的要求，我们可以构建一个使用SMA或EMA收敛与发散概念的策略，例如利用短期和长期均线的交叉来模拟类似MACD的交易逻辑。以下是一个基于SMA收敛发散思想的简单策略示例：
#这个策略中，我们定义了短期（short_period）和长期（long_period）均线周期，当短期均线向上穿越长期均线时，视为买入信号；当短期均线向下穿越长期均线时，视为卖出信号。这实质上是利用均线收敛（交叉）作为市场趋势变化的信号，类似于MACD策略中利用快线和慢线的交叉来判断买卖时机。